﻿# -*- coding: utf-8 -*-
import requests
import chromadb
import uuid
from flask import Flask, request, jsonify
from flask_cors import CORS
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建Flask应用
app = Flask(__name__)
CORS(app)

# 配置
EMBEDDING_API_URL = "http://172.16.35.50:8000/v1/embeddings"
CHROMA_HOST = "172.16.29.101"
CHROMA_PORT = 23000

# 全局变量
chroma_client = None
collection = None

def init_chroma():
    global chroma_client, collection
    try:
        logger.info(f"连接到Chroma服务器: {CHROMA_HOST}:{CHROMA_PORT}")
        chroma_client = chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)
        version_info = chroma_client.get_version()
        logger.info(f"Chroma连接成功！版本: {version_info}")
        collection = chroma_client.get_or_create_collection(name="technical_standards")
        logger.info("集合已准备就绪")
        return True
    except Exception as e:
        logger.error(f"Chroma连接失败: {str(e)}")
        return False

def get_embedding(text):
    try:
        payload = {"model": "bge-large-zh-v1.5", "input": text}
        response = requests.post(EMBEDDING_API_URL, json=payload)
        response.raise_for_status()
        data = response.json()
        return data["data"][0]["embedding"]
    except Exception as e:
        logger.error(f"获取嵌入向量失败: {str(e)}")
        return None

@app.route("/api/health", methods=["GET"])
def health_check():
    try:
        if chroma_client is None or collection is None:
            return jsonify({"status": "error", "message": "Chroma服务未初始化"}), 500
        version_info = chroma_client.get_version()
        collection_count = collection.count()
        return jsonify({
            "status": "success",
            "message": "服务正常运行",
            "chroma_version": version_info,
            "collection_count": collection_count
        })
    except Exception as e:
        return jsonify({"status": "error", "message": f"健康检查失败: {str(e)}"}), 500

@app.route("/api/search", methods=["POST"])
def search_documents():
    try:
        data = request.get_json()
        if not data or "query" not in data:
            return jsonify({"status": "error", "message": "需要包含query字段"}), 400

        query_text = data["query"]
        n_results = data.get("n_results", 3)

        query_vector = get_embedding(query_text)
        if query_vector is None:
            return jsonify({"status": "error", "message": "无法生成查询向量"}), 500

        results = collection.query(
            query_embeddings=[query_vector],
            n_results=n_results
        )

        search_results = []
        if results and results["ids"][0]:
            for i in range(len(results["ids"][0])):
                result_item = {
                    "id": results["ids"][0][i],
                    "document": results["documents"][0][i],
                    "distance": results["distances"][0][i] if "distances" in results and results["distances"][0] else None,
                    "metadata": results["metadatas"][0][i] if "metadatas" in results and results["metadatas"][0] else {}
                }
                search_results.append(result_item)

        logger.info(f"查询完成，找到 {len(search_results)} 个结果")
        return jsonify({
            "status": "success",
            "message": "查询完成",
            "query": query_text,
            "results": search_results,
            "total_results": len(search_results)
        })
    except Exception as e:
        logger.error(f"查询失败: {str(e)}")
        return jsonify({"status": "error", "message": f"查询失败: {str(e)}"}), 500

if __name__ == "__main__":
    if not init_chroma():
        logger.error("Chroma初始化失败，退出应用")
        exit(1)
    
    logger.info("Chroma API服务启动中...")
    logger.info("API接口说明:")
    logger.info("  GET  /api/health - 健康检查")
    logger.info("  POST /api/search - 搜索相似文档")
    logger.info("服务地址: http://0.0.0.0:24000")
    
    app.run(host="172.16.29.101", port=24000, debug=True)
